#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <libspe2.h>
#include <math.h>
#include <pthread.h>

#include "fish.h"
#include "flocking.h"

#define VERSION 2

/* Flocking Parameters */
const double C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const double C1 = 2000.0; // weight of vector towards center of world
const double C1 = 0.0;
const double C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const double C3 = 0.0; // weight of vector towards center of gravity
const double C4 = 0.05; // weight of random vector
const double C5 = 0.1; // weight of vector pointing to cog of local flockmates
const double R_CROWD2 = 0.00390625; // square of crowding radius
const double R_SIGHT2 = 0.015625; // square of sighting radius
const double R_SIGHT = 0.125; // sighting radius
const double BIRD_MASS = 1.0; // bird mass
const double R = 1.0; // repulsion constant
const double M = 1.0; // constant that converts vectors to acceleration
const double START_CONST = 0.5;
const double SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const double SCALE_FACTOR = 1.0;
double v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay; // vectors
double cogx, cogy; // center of gravity coordinates
double local_cogx, local_cogy;
double tx, ty; // translation vectors for the moving grid


FILE *output_fp; /* Pointer to output file */
FILE *input_fp; /* Pointer to input file */
int outputp;
int inputp;

int n_fish = 10; /* Other defaults in fish-lib.c. */
//fish_t *fish; /* Global array of fish */

#define BUCKETS 2

struct Timer total_timer;
struct Timer move_timer;
struct Timer compute_dt_timer;
struct Timer interact_fish_timer;
struct Timer spe_program_load_timer;
struct Timer spe_setup_timer;
struct Timer pthreads_run_timer;
struct Timer spe_context_run_timer[BUCKETS][BUCKETS];

fish_t* sorted;
fish_t *fish[2][BUCKETS][BUCKETS];    /* Buckets of fish */
int fish_bucket_size[BUCKETS][BUCKETS]; /* Capacities of buckets */
int fish_bucket_n[2][BUCKETS][BUCKETS]; /* How many fish in each bucket */


/* SPE PROGRAM STUFF */
spe_params_t spe_params __attribute__((aligned(16)));
spe_auxinfo_t spe_auxinfo[BUCKETS][BUCKETS] __attribute__((aligned(16)));
spe_context_ptr_t spe_contexts[BUCKETS][BUCKETS];
spe_program_handle_t *prog;
pthread_t threads[BUCKETS][BUCKETS];
double dt;
void* run_sim_spe(void*);

typedef struct {
  spe_context_ptr_t spe;
  int bucket_number;
} thread_args_t;
thread_args_t thread_args[BUCKETS][BUCKETS];  

struct Timer total_timer;

void init_fish();
void interact_fish(int t);
void interact_fish_bucket(int t, int n, int m);
void bounce_fish (fish_t* fish);
void move_fish (double dt, int t);
double compute_norm (int t);
void output_fish(FILE* output_fp, const double curr_time,
    const double dt, const int t);
int fish_bucket_x(fish_t const* f);
int fish_bucket_y(fish_t const* f);
void sort_fish(int);


int
main(int argc, char **argv)
{
  double curr_time;
  double output_time;
  int t, steps;

  dt = 0.0;

  get_options(argc, argv);
  //srand48(clock());

  if (output_filename) {
    outputp = 1;
    output_fp = fopen(output_filename, "w");
    if (output_fp == NULL) {
      printf("Could not open %s for output\n", output_filename);
      exit(1);
    }
    fprintf(output_fp, "n_fish: %d\n", n_fish);
  }

  if (input_filename) {
    inputp = 1;
    input_fp = fopen(input_filename, "r");
    if (input_fp == NULL) {
      printf("Could not open %s for input\n", input_filename);
      exit(1);
    }
  }

  init_fish();

  // SPE STUFF
  prog = spe_image_open("./fish_spe.elf");
  if( !prog ) {
    perror("spe_image_open");
    exit(1);
  }

  
  spe_auxinfo[0][0].bucket_number = 0;
  spe_auxinfo[0][1].bucket_number = 1;
  spe_auxinfo[1][0].bucket_number = 2;
  spe_auxinfo[1][1].bucket_number = 3;

  if (inputp) { 
    fclose(input_fp);
  }
  

  initialize_timer (&total_timer);
  initialize_timer (&move_timer);
  initialize_timer (&compute_dt_timer);
  initialize_timer (&interact_fish_timer);
  initialize_timer (&spe_program_load_timer);
  initialize_timer (&spe_setup_timer);
  initialize_timer (&pthreads_run_timer);
  initialize_timer (&(spe_context_run_timer[0][0]));
  initialize_timer (&(spe_context_run_timer[0][1]));
  initialize_timer (&(spe_context_run_timer[1][0]));
  initialize_timer (&(spe_context_run_timer[1][1]));

  start_timer(&total_timer);

  for (output_time = 0.0, curr_time = 0.0, steps = 0, t = 0;
       curr_time <= end_time && steps < max_steps;
       curr_time += dt, ++steps, t ^= 1) {

    /* We only output once every output_interval time unit, at most.
     * Without that restriction, we can easily create a huge output
     * file.  Printing a record for ten fish takes about 300 bytes,
     * so for every 1000 steps, we could dump 300K of info.  Now
     * scale the number of fish by 1000...
     */

  
    if (outputp && curr_time >= output_time) {
      output_fish(output_fp, curr_time, dt, t);
      output_time = curr_time + output_interval;
    }
  
    /* generally what we want to do 
     * interact_fish(t)
     *  - spe's get all the fish
     *  - spe's calculate the accellerations on the fish in their bucket
     *  - spe's reduce max dt
     *  - spe's move their fish, but keep the fish in their buckets
     *  - spe's write back their buckets to main memory
     * sort_fish(t)
     *  - resort the fish into new buckets
     */

    // here is where we invoke the spes to do some parallel work for us
    //printf("calling interact_fish()\n");
    start_timer(&interact_fish_timer);
    interact_fish(t);
    stop_timer(&interact_fish_timer);
    
    //printf("interact_fish worked!\n");

    start_timer(&move_timer);
    sort_fish(t);
    stop_timer(&move_timer);
    //printf("sort fish worked!\n");
  }

  stop_timer(&total_timer);

  if (outputp) {
    output_fish(output_fp, curr_time, dt, t);
    fclose(output_fp);
  }

  printf("version\t%d\tfish\t%d\tcurr_time\t%g\tend_time\t%g\tsteps\t%d\tmax_steps\t%d\t", VERSION, n_fish, curr_time, end_time, steps, max_steps );
  printf("total_timer\t%g\t", timer_duration(total_timer));
  printf("sort_timer\t%g\t", timer_duration(move_timer));
  printf("compute_dt_timer\t%g\t", timer_duration(compute_dt_timer));
  printf("interact_fish_timer\t%g\t", timer_duration(interact_fish_timer));
  printf("spe_program_load_timer\t%g\t", timer_duration(spe_program_load_timer));
  printf("spe_setup_timer\t%g\t", timer_duration(spe_setup_timer));
  printf("pthreads_run_timer\t%g\t", timer_duration(pthreads_run_timer));
  printf("spe_context_run_timer[0][0]\t%g\t", timer_duration(spe_context_run_timer[0][0]));
  printf("spe_context_run_timer[0][1]\t%g\t", timer_duration(spe_context_run_timer[0][1]));
  printf("spe_context_run_timer[1][0]\t%g\t", timer_duration(spe_context_run_timer[1][0]));
  printf("spe_context_run_timer[1][1]\t%g\t", timer_duration(spe_context_run_timer[1][1]));
  printf("\n");


  // first number is the simulation version                                                                                                                          
  //  printf("%d\t%d\t%g\t%g\t%d\t%d\t%g\n", 2, n_fish,
  //curr_time, end_time, steps, max_steps, timer_duration(total_timer));
  /*
  printf("Ended at %g (%g), %d (%d) steps\n",
      curr_time, end_time, steps, max_steps);
  printf("Total time taken: %g\n", timer_duration(total_timer));
  */
  return 0;
}


void
sort_fish(int t)
{
  int x, y, i, nx, ny, tp1;
  
  tp1 = t^1;

  cogx = cogy = 0.0;

  //printf("memsetting\n");
  memset( fish_bucket_n[tp1], 0, BUCKETS*BUCKETS*sizeof(int) );

  for(x = 0; x < BUCKETS; x++){
    for(y = 0; y < BUCKETS; y++){
      for(i = 0; i < fish_bucket_n[t][x][y]; i++)
      {
	//printf("accessing a fish\n");
	fish_t *__restrict__ f = &fish[t][x][y][i];

	cogx += f->x;
	cogy += f->y;

        nx = fish_bucket_x(f);
        ny = fish_bucket_y(f);

	//printf("moving a fish\n");
	//printf("details:\n\tnx: %d\n\tny: %d\n\ttp1: %d\n\tfish_bucket_n[tp1][nx][ny]: %d\n", nx, ny, tp1, fish_bucket_n[tp1][nx][ny]);
        fish[tp1][nx][ny][fish_bucket_n[tp1][nx][ny]++]
          = fish[t][x][y][i]; 
      }
    }
  }

  //printf("out of the loop\n");

  cogx /= n_fish;
  cogy /= n_fish;
  //printf("calculated cogx/cogy\n");
  tx = scale_coord(cogx) - 0.5;
  ty = scale_coord(cogy) - 0.5;
  //printf("finished sort_fish()\n");

}


/* OK so what we are going to do here is give the 
 * spes a pointer to the information that they need
 * this information includes a pointer to their bucket
 * an array of pointers to their neighbor buckets
 * and 4 integers that represent the # of fish in 
 * each bucket. I'm assuming that we're dealing with
 * a 2x2 grid, so each bucket has 3 true neighbor
 * buckets. 
 */

void
interact_fish(int t)
{
  //printf("starting interact_fish()\n");
  int bucket_x, bucket_y, ret;
  double max_norm;
  
  start_timer(&spe_setup_timer);

  // perhaps should be typecasted to unsigned long long? - IBM
  // suggests this
  spe_params.nw = (unsigned long)fish[t][0][0];
  spe_params.ne = (unsigned long)fish[t][0][1];
  spe_params.sw = (unsigned long)fish[t][1][0];
  spe_params.se = (unsigned long)fish[t][1][1];
  spe_params.nw_n = fish_bucket_n[t][0][0];
  spe_params.ne_n = fish_bucket_n[t][0][1];
  spe_params.sw_n = fish_bucket_n[t][1][0];
  spe_params.se_n = fish_bucket_n[t][1][1];

  spe_auxinfo[0][0].random_seed = drand48();
  spe_auxinfo[0][1].random_seed = drand48();
  spe_auxinfo[1][0].random_seed = drand48();
  spe_auxinfo[1][1].random_seed = drand48();

  //printf("creating contexts and loading programs\n");

  for(bucket_x = 0; bucket_x < BUCKETS; bucket_x++){
    for(bucket_y = 0; bucket_y < BUCKETS; bucket_y++){
      thread_args[bucket_x][bucket_y].spe = spe_context_create(0, NULL);
      if( !thread_args[bucket_x][bucket_y].spe ) {
	perror("spe_context_create");
	exit(1);
      }
      start_timer(&spe_program_load_timer);
      ret = spe_program_load( thread_args[bucket_x][bucket_y].spe, prog );
      stop_timer(&spe_program_load_timer);
      if( ret ) {
	perror("spe_program_load");
	exit(1);
      }

      thread_args[bucket_x][bucket_y].bucket_number = bucket_x*BUCKETS + bucket_y;
    }
  }

  stop_timer(&spe_setup_timer);

  //printf("creating pthreads\n");
  start_timer(&pthreads_run_timer);

  for(bucket_x = 0; bucket_x < BUCKETS; bucket_x++){
    for(bucket_y = 0; bucket_y < BUCKETS; bucket_y++){
      ret = pthread_create( &threads[bucket_x][bucket_y], NULL, run_sim_spe, &thread_args[bucket_x][bucket_y] );
      if( ret ) {
	perror("pthread_create");
	exit(1);
      }
    }
  }

  // Before joining on the threads we need to be waiting to help out w/ the reduce operation
  // The dt reduce op. will go like this:
  //  - poll spe outbound mailboxes for messages (these will be their max_norm's)
  //  - gather all the max_norm's
  //  - find the max of all of them
  //  - compute dt
  //  - send back dt.

  

  max_norm = 0.0;

  union {
    float f;
    int i;
  } u;

  //printf("attempting to get the max_norms from the spes\n");

  int flag = 0;

  for( bucket_x = 0; bucket_x < BUCKETS; bucket_x++ ){
    for( bucket_y = 0; bucket_y < BUCKETS; bucket_y++ ){
      while( !spe_out_mbox_status( thread_args[bucket_x][bucket_y].spe ) ) {
	//printf("%d", bucket_x*BUCKETS + bucket_y);
      }
      if( !flag ) {
	start_timer(&compute_dt_timer);
	flag = 1;
      }

      //printf("attempting to read from %d's mailbox\n", bucket_x*BUCKETS + bucket_y);
      spe_out_mbox_read( thread_args[bucket_x][bucket_y].spe, &(u.i), 1); // tee hee!
      //printf("Got %f from %d\n", u.f, bucket_x*BUCKETS + bucket_y);
      max_norm = f_max(max_norm, (double)(u.f));
    }
  }

  //printf("got all the max_norms\n");
  
  dt = max_norm_change / max_norm;
  dt = f_max(dt, min_dt);
  dt = f_min(dt, max_dt);
  u.f = (float)dt;
  dt = (double)(u.f); // this is to eliminate any discrepancies between the ppe's and spe's dt value

  unsigned int msg0 = u.i;
  unsigned int msg1 = u.i;
  unsigned int msg2 = u.i;
  unsigned int msg3 = u.i;

  //printf("Sending dt: %f\n", dt);

  while( !spe_in_mbox_status( thread_args[0][0].spe ) ) {}
  spe_in_mbox_write( thread_args[0][0].spe, &msg0, 1, SPE_MBOX_ALL_BLOCKING );

  while( !spe_in_mbox_status( thread_args[0][1].spe ) ) {}
  spe_in_mbox_write( thread_args[0][1].spe, &msg1, 1, SPE_MBOX_ALL_BLOCKING );

  while( !spe_in_mbox_status( thread_args[1][0].spe ) ) {}
  spe_in_mbox_write( thread_args[1][0].spe, &msg2, 1, SPE_MBOX_ALL_BLOCKING );

  while( !spe_in_mbox_status( thread_args[1][1].spe ) ) {}
  spe_in_mbox_write( thread_args[1][1].spe, &msg3, 1, SPE_MBOX_ALL_BLOCKING );

  stop_timer(&compute_dt_timer);

  //printf("got here!\n");

  /*
  // send it back
  for( bucket_x = 0; bucket_x < BUCKETS; bucket_x++ ) {
    for( bucket_y = 0; bucket_y < BUCKETS; bucket_y++ ) {
      if( -1 == spe_in_mbox_write( thread_args[bucket_x][bucket_y].spe, &msg, 1, SPE_MBOX_ALL_BLOCKING ) ){
	perror("spe_in_mbox_write()");
	exit(1);
      }
    }
  }
  */
 

  for(bucket_x = 0; bucket_x < BUCKETS; bucket_x++) {
    for(bucket_y = 0; bucket_y < BUCKETS; bucket_y++) {
      //printf("joining on: %d\n", bucket_x*BUCKETS + bucket_y);
      pthread_join( threads[bucket_x][bucket_y], NULL );
      //printf("thread %d finished\n", bucket_x*BUCKETS + bucket_y);
      ret = spe_context_destroy(thread_args[bucket_x][bucket_y].spe);
      //printf("destroyed context %d\n", bucket_x*BUCKETS + bucket_y );
      if( ret ) {
	//printf("ret != 0\n");
	perror("spe_context_destroy");
	exit(1);
      }
    }
  }

  stop_timer(&pthreads_run_timer);

  //printf("I'm at the end of my rope!\n");
}


/* Thread function to run our program on an spe
 */
void*
run_sim_spe(void* thread_arg)
{
  int ret;
  thread_args_t* args = (thread_args_t*)thread_arg;
  unsigned int entry;
  spe_stop_info_t stop_info;

  int n = (args->bucket_number)/BUCKETS;
  int m = (args->bucket_number)%BUCKETS;

  entry = SPE_DEFAULT_ENTRY;
  //printf("Starting up context: %d\n", n*BUCKETS + m);
  start_timer(&spe_context_run_timer[n][m]);
  ret = spe_context_run( args->spe, &entry, 0, &spe_params, &spe_auxinfo[n][m], &stop_info );
  stop_timer(&spe_context_run_timer[n][m]);
  //printf("Context %d finished\n", n*BUCKETS + m);
  if( ret < 0 ) {
    perror("spe_context_run");
    return NULL;
  }
  

  return NULL;
}


/* Allocate and initialize the fish positions / velocities / accelerations. */

void
init_fish()
{
  int i, j, nx, ny;
  float xf, yf;
  fish_t *fishes;

  /* I'll just do this the easy way - allocate and initialize them all at
    once, and then split them into buckets. This way is easier to code, and
    results in fewer calls into stdlib malloc/realloc etc. */

  if(!(fishes = (fish_t*)malloc(n_fish*sizeof(fish_t)))){
    perror("malloc");
    exit(-1);
  }
  
  if(uniformp){
    /*srand48(clock());*/
    srand48(0);
  }
  
  cogx = cogy = 0;

  for(i = 0; i < n_fish; i++){
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].x = unscale_coord(xf);
      fishes[i].y = unscale_coord(yf);
    }else if(uniformp){
      fishes[i].x = unscale_coord(drand48());
      fishes[i].y = unscale_coord(drand48());
    }else{
      double angle = i * (2.0 * M_PI / n_fish);
      fishes[i].x = unscale_coord(0.5 * cos(angle) + 0.5);
      fishes[i].y = unscale_coord(0.5 * sin(angle) + 0.5);
    }
    cogx += fishes[i].x;
    cogy += fishes[i].y;
  }

  cogx /= n_fish;
  cogy /= n_fish;
  tx = scale_coord(cogx) - 0.5;
  ty = scale_coord(cogy) - 0.5;

  for(i = 0; i < n_fish; i++){
    fishes[i].ax = fishes[i].ay = 0.0;
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].vx = xf;
      fishes[i].vy = yf;
    }else{
      fishes[i].vx = START_CONST*(drand48() - 0.5);
      fishes[i].vy = START_CONST*(drand48() - 0.5);
    }
  }

  for(i = 0; i < BUCKETS; i++){
    for(j = 0; j < BUCKETS; j++)
    {
      fish_bucket_size[i][j] = 1;
      fish_bucket_n[0][i][j] = 0;
      fish_bucket_n[1][i][j] = 0;
    }
  }

  for(i = 0; i < n_fish; i++){
    nx = fish_bucket_x(&fishes[i]);
    ny = fish_bucket_y(&fishes[i]);
    fish_bucket_size[nx][ny]++; 
  }

  for(i = 0; i < BUCKETS; i++){
    for(j = 0; j < BUCKETS; j++) {
      // int size = sizeof(fish_t)*fish_bucket_size[i][j];
      // use memalign to allocate an aligned piece of
      // memory... to 16 bytes for DMA transfers
      int size = sizeof(fish_t)*n_fish;
      if(!(fish[0][i][j] = (fish_t*)memalign(16, size)) ||
	 !(fish[1][i][j] = (fish_t*)memalign(16, size)))
      {
        perror("memalign");
        exit(-1);
      }
    }
  }

  for(i = 0; i < n_fish; i++){
    nx = fish_bucket_x(&fishes[i]);
    ny = fish_bucket_y(&fishes[i]);
    fish[0][nx][ny][ fish_bucket_n[0][nx][ny]++ ] = fishes[i];
  }

  free(fishes);

  sorted = malloc(n_fish*sizeof(fish_t));
  if(!sorted){
    perror("malloc");
    exit(-1);
  }
}

/*
  Dump out all the fishies (and their center of gravity)
  in a format that the viewer understands.
*/


void
output_fish(FILE* output_fp, const double curr_time,
    const double dt, const int t)
{
  int i, j, x, y;

  j = 0;
  for(x = 0; x < BUCKETS; x++){
    for(y = 0; y < BUCKETS; y++){
      for(i = 0; i < fish_bucket_n[t][x][y]; i++){
        sorted[j] = fish[t][x][y][i];
	j++;
      }
    }
  }

  fprintf(output_fp, "%.5g (%.5g):\n", curr_time, dt);
  for (i = 0; i < n_fish; ++i) {
    fprintf(output_fp, "  %d: (%g, %g)\n", i,
        scale_coord(sorted[i].x), scale_coord(sorted[i].y));
  }
  if (outvelp) {
    for (i =0; i< n_fish; ++i) {
      fprintf(output_fp, "  %d: (%g, %g)\n", i, sorted[i].vx, sorted[i].vy);
    }
  }
  fprintf(output_fp, "  cg: (%g, %g)\n", scale_coord(cogx),
      scale_coord(cogy));
}

int
fish_bucket_x(fish_t const* f)
{
  //return floor(fmod(scale_coord(f->x)-tx+1,1)*BUCKETS); 
  return floor(fmod(scale_coord(f->x)-tx+1,1)*BUCKETS); 
}

int
fish_bucket_y(fish_t const* f)
{
  //return floor(fmod(scale_coord(f->y)-ty+1,1)*BUCKETS); 
  return floor(fmod(scale_coord(f->y)-ty+1,1)*BUCKETS); 
}


